from torch import nn
from cassava_leaf.utils import get_model
from fastai.vision.all import create_head, LabelSmoothingCrossEntropy, \
    Callback, Metric, accuracy, find_bs
import torch.nn.functional as F


class MultiHeadModel(nn.Module):

    def __init__(self, backbone_name):
        super(MultiHeadModel, self).__init__()
        self.feature = get_model(backbone_name)
        if hasattr(self.feature, 'fc'):
            self.last_feat_dim = self.feature.fc.in_features
            self.feature.fc = nn.Identity()
        elif hasattr(self.feature, 'classifier'):
            self.last_feat_dim = self.feature.classifier.in_features
            self.feature.classifier = nn.Identity()
        else:
            raise Exception()
        self.feature.global_pool = nn.Identity()
        # self.head1 = nn.Linear(self.last_feat_dim, 5)
        # self.head2 = nn.Linear(self.last_feat_dim, 2)
        self.head1 = create_head(self.last_feat_dim * 2, 5)
        self.head2 = create_head(self.last_feat_dim * 2, 2)

    def forward(self, x):
        feature = self.feature(x)
        return self.head1(feature), self.head2(feature)


class MultiOutputLoss(nn.Module):

    def __init__(self, use_label_smoothing, eps):
        super(MultiOutputLoss, self).__init__()
        self.use_label_smoothing = use_label_smoothing
        if use_label_smoothing:
            self.loss1 = LabelSmoothingCrossEntropy(eps)
        else:
            self.loss1 = F.cross_entropy

    def forward(self, y_pred, y_true_0, y_true_1):
        return self.loss1(y_pred[0], y_true_0) + F.cross_entropy(y_pred[1], y_true_1)


class MultiOutputCallback(Callback):

    def before_batch(self):
        old_yb = self.learn.yb[0]
        bin_yb = old_yb.clone()
        bin_yb[bin_yb != 4] = 0
        bin_yb[bin_yb == 4] = 1
        self.learn.yb = (old_yb, bin_yb)
        # print(self.learn.yb)


class AccMul(Metric):

    def __init__(self):
        self.func = accuracy

    def reset(self):
        self.total, self.count = 0., 0

    def accumulate(self, learn):
        bs = find_bs(learn.yb)
        self.total += learn.to_detach(self.func(learn.pred[0], learn.yb[0]))*bs
        self.count += bs

    @property
    def value(self):
        return self.total/self.count if self.count != 0 else None

    @property
    def name(self):
        return 'AccMul'


class AccBin(Metric):

    def __init__(self):
        self.func = accuracy

    def reset(self):
        self.total, self.count = 0., 0

    def accumulate(self, learn):
        bs = find_bs(learn.yb)
        self.total += learn.to_detach(self.func(learn.pred[1], learn.yb[1]))*bs
        self.count += bs

    @property
    def value(self):
        return self.total/self.count if self.count != 0 else None

    @property
    def name(self):
        return 'AccBin'
